import cv2
import numpy as np

# 读取视频
cap = cv2.VideoCapture('test.mp4')

# 检查视频是否成功打开
if not cap.isOpened():
    print("Error opening video file")
    exit()

# 创建窗口和滑动条
cv2.namedWindow('Red Line Detection')

# 定义红色HSV阈值范围（红色在HSV中分为两部分）
def nothing(x):
    pass

# 创建滑动条
cv2.createTrackbar('HMin1', 'Red Line Detection', 0, 179, nothing)
cv2.createTrackbar('SMin1', 'Red Line Detection', 100, 255, nothing)
cv2.createTrackbar('VMin1', 'Red Line Detection', 100, 255, nothing)
cv2.createTrackbar('HMax1', 'Red Line Detection', 15, 179, nothing)
cv2.createTrackbar('SMax1', 'Red Line Detection', 255, 255, nothing)
cv2.createTrackbar('VMax1', 'Red Line Detection', 255, 255, nothing)

cv2.createTrackbar('HMin2', 'Red Line Detection', 165, 179, nothing)
cv2.createTrackbar('SMin2', 'Red Line Detection', 100, 255, nothing)
cv2.createTrackbar('VMin2', 'Red Line Detection', 100, 255, nothing)
cv2.createTrackbar('HMax2', 'Red Line Detection', 179, 179, nothing)
cv2.createTrackbar('SMax2', 'Red Line Detection', 255, 255, nothing)
cv2.createTrackbar('VMax2', 'Red Line Detection', 255, 255, nothing)

# 读取第一帧作为初始帧
ret, frame = cap.read()
if not ret:
    print("Error reading video frame")
    exit()

# 当前帧号
frame_count = 0

while True:
    # 获取当前滑动条值
    hmin1 = cv2.getTrackbarPos('HMin1', 'Red Line Detection')
    smin1 = cv2.getTrackbarPos('SMin1', 'Red Line Detection')
    vmin1 = cv2.getTrackbarPos('VMin1', 'Red Line Detection')
    hmax1 = cv2.getTrackbarPos('HMax1', 'Red Line Detection')
    smax1 = cv2.getTrackbarPos('SMax1', 'Red Line Detection')
    vmax1 = cv2.getTrackbarPos('VMax1', 'Red Line Detection')
    
    hmin2 = cv2.getTrackbarPos('HMin2', 'Red Line Detection')
    smin2 = cv2.getTrackbarPos('SMin2', 'Red Line Detection')
    vmin2 = cv2.getTrackbarPos('VMin2', 'Red Line Detection')
    hmax2 = cv2.getTrackbarPos('HMax2', 'Red Line Detection')
    smax2 = cv2.getTrackbarPos('SMax2', 'Red Line Detection')
    vmax2 = cv2.getTrackbarPos('VMax2', 'Red Line Detection')
    
    # 转换为HSV颜色空间
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    
    # 创建红色掩码（两部分）
    lower_red1 = np.array([hmin1, smin1, vmin1])
    upper_red1 = np.array([hmax1, smax1, vmax1])
    lower_red2 = np.array([hmin2, smin2, vmin2])
    upper_red2 = np.array([hmax2, smax2, vmax2])
    
    mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
    
    # 合并两个掩码
    mask = cv2.bitwise_or(mask1, mask2)
    
    # 对原始图像和掩码进行位运算
    result = cv2.bitwise_and(frame, frame, mask=mask)

    binary_mask = np.zeros((result.shape[0], result.shape[1]), dtype=np.uint8)
    non_black_pixels = np.any(result > 0, axis=-1)
    binary_mask[non_black_pixels] = 255
    
    # 在结果图像上显示当前阈值
    threshold_text = f'H1: {hmin1}-{hmax1}, S1: {smin1}-{smax1}, V1: {vmin1}-{vmax1}'
    threshold_text2 = f'H2: {hmin2}-{hmax2}, S2: {smin2}-{smax2}, V2: {vmin2}-{vmax2}'
    frame_text = f'Frame: {frame_count}'
    
    cv2.putText(result, threshold_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    cv2.putText(result, threshold_text2, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    cv2.putText(result, frame_text, (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    
    # 显示结果
    cv2.imshow('Original', frame)
    cv2.imshow('Red Line Detection', result)
    print(result)
    
    # 保存当前帧的阈值和结果
    key = cv2.waitKey(0) & 0xFF
    
    # 空格键：下一帧
    if key == 32:  # 空格键
        ret, frame = cap.read()
        if not ret:
            print("End of video")
            break
        frame_count += 1
        print(f"Current frame: {frame_count}")
    
    # 's'键：保存当前结果和阈值
    elif key == ord('s'):
        result_filename = f'result_frame_{frame_count}.jpg'
        cv2.imwrite(result_filename, result)
        print(f"保存结果到: {result_filename}")
        
        threshold_filename = f'threshold_frame_{frame_count}.txt'
        with open(threshold_filename, 'w') as f:
            f.write(f'H1: {hmin1}-{hmax1}, S1: {smin1}-{smax1}, V1: {vmin1}-{vmax1}\n')
            f.write(f'H2: {hmin2}-{hmax2}, S2: {smin2}-{smax2}, V2: {vmin2}-{vmax2}\n')
        print(f"保存阈值到: {threshold_filename}")
    
    # 'q'键：退出程序
    elif key == ord('q'):
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()